{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3D Multi-organ Segmentation with Swin UNETR\n",
    "This tutorial uses a Swin UNETR [1] model for the task of multi-organ segmentation task using the BTCV challenge dataset. The architecture of Swin UNETR is demonstrated as below\n",
    "\n",
    "![image](../figures/swin_unetr_btcv.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup environment\n",
    "`pip install  \"monai-weekly[nibabel, tqdm]\" -i https://pypi.tuna.tsinghua.edu.cn/simple/` <br/>\n",
    "`pip install  matplotlib - i https://pypi.tuna.tsinghua.edu.cn/simple/` <br/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MONAI version: 1.4.dev2426\n",
      "Numpy version: 1.26.4\n",
      "Pytorch version: 2.3.1+cu121\n",
      "MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False\n",
      "MONAI rev id: d622a16f927841fdd7d057b7553805405f0805e4\n",
      "MONAI __file__: /home/<username>/anaconda3/envs/monai/lib/python3.9/site-packages/monai/__init__.py\n",
      "\n",
      "Optional dependencies:\n",
      "Pytorch Ignite version: 0.4.11\n",
      "ITK version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "Nibabel version: 5.2.1\n",
      "scikit-image version: 0.24.0\n",
      "scipy version: 1.13.1\n",
      "Pillow version: 10.4.0\n",
      "Tensorboard version: 2.17.0\n",
      "gdown version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "TorchVision version: 0.18.1+cu121\n",
      "tqdm version: 4.66.4\n",
      "lmdb version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "psutil version: 6.0.0\n",
      "pandas version: 2.2.2\n",
      "einops version: 0.8.0\n",
      "transformers version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "mlflow version: 2.14.2\n",
      "pynrrd version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "clearml version: NOT INSTALLED or UNKNOWN VERSION.\n",
      "\n",
      "For details about installing the optional dependencies, please visit:\n",
      "    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import shutil\n",
    "import tempfile\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "\n",
    "from monai.losses import DiceCELoss\n",
    "from monai.inferers import sliding_window_inference\n",
    "from monai.transforms import (\n",
    "    AsDiscrete,\n",
    "    Compose,\n",
    "    CropForegroundd,\n",
    "    LoadImaged,\n",
    "    Orientationd,\n",
    "    RandFlipd,\n",
    "    RandCropByPosNegLabeld,\n",
    "    RandShiftIntensityd,\n",
    "    ScaleIntensityRanged,\n",
    "    Spacingd,\n",
    "    RandRotate90d,\n",
    "    EnsureTyped,\n",
    ")\n",
    "\n",
    "from monai.config import print_config\n",
    "from monai.metrics import DiceMetric\n",
    "from monai.networks.nets import SwinUNETR\n",
    "\n",
    "from monai.data import (\n",
    "    ThreadDataLoader,\n",
    "    CacheDataset,\n",
    "    load_decathlon_datalist,\n",
    "    decollate_batch,\n",
    "    set_track_meta,\n",
    ")\n",
    "import torch\n",
    "\n",
    "print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup data directory\n",
    "\n",
    "You can specify a directory with the `MONAI_DATA_DIRECTORY` environment variable.  \n",
    "This allows you to save results and reuse downloads.  \n",
    "If not specified a temporary directory will be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/media/chyang/data/dataset/project/registration/hipbone_raw_data/train/\n"
     ]
    }
   ],
   "source": [
    "os.environ['MONAI_DATA_DIRECTORY'] = '/media/chyang/data/dataset/project/registration/hipbone_raw_data/train/'\n",
    "directory = os.environ.get(\"MONAI_DATA_DIRECTORY\")\n",
    "if directory is not None:\n",
    "    os.makedirs(directory, exist_ok=True)\n",
    "root_dir = tempfile.mkdtemp() if directory is None else directory\n",
    "print(root_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup transforms for training and validation\n",
    "To save on GPU memory utilization, the num_samples can be reduced to 2. \n",
    "\n",
    "A note on design related to MetaTensors:\n",
    "\n",
    "- Summary: using `EnsureTyped(..., track_meta=False)` (caching) and `set_track_meta(False)` (during training) speeds up training significantly.\n",
    "\n",
    "- We are moving towards the use of MONAI's MetaTensor in place of numpy arrays or PyTorch tensors. MetaTensors have the benefit of carrying the metadata directly with the tensor, but in some use cases (like here with training, where training data are only used for computing loss and metadata is not useful), we can safely disregard the metadata to improve speed.\n",
    "\n",
    "- Hence, you will see `EnsureTyped` being used before the first random transform in the training transform chain, which caches the result of deterministic transforms on GPU as Tensors (rather than MetaTensors), with `track_meta = False`. \n",
    "\n",
    "- On the other hand, in the following demos we will display example validation images, which uses metadata, so we use `EnsureTyped` with `track_meta = True`. Since there are no random transforms during validation, tracking metadata for validation images causes virtually no slowdown (~0.5%).\n",
    "\n",
    "- In the next section, you will see `set_track_meta(False)`. This is a global API introduced in MONAI 0.9.1, and it makes sure that random transforms will also be performed using Tensors rather than MetaTensors. Used together with `track_meta=False` in `EnsureTyped`, it results in all transforms being performed on Tensors, which we have found to speed up training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 4\n",
    "\n",
    "# os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "# device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "train_transforms = Compose(\n",
    "    [\n",
    "        LoadImaged(keys=[\"image\", \"label\"], ensure_channel_first=True),\n",
    "        ScaleIntensityRanged(\n",
    "            keys=[\"image\"],\n",
    "            a_min=-175,\n",
    "            a_max=250,\n",
    "            b_min=0.0,\n",
    "            b_max=1.0,\n",
    "            clip=True,\n",
    "        ),\n",
    "        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
    "        Spacingd(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            pixdim=(1.5, 1.5, 2.0),\n",
    "            mode=(\"bilinear\", \"nearest\"),\n",
    "        ),\n",
    "        EnsureTyped(keys=[\"image\", \"label\"], device=device, track_meta=False),\n",
    "        RandCropByPosNegLabeld(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            label_key=\"label\",\n",
    "            spatial_size=(96, 96, 96),\n",
    "            pos=1,\n",
    "            neg=1,\n",
    "            num_samples=num_samples,\n",
    "            image_key=\"image\",\n",
    "            image_threshold=0,\n",
    "        ),\n",
    "        RandFlipd(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            spatial_axis=[0],\n",
    "            prob=0.10,\n",
    "        ),\n",
    "        RandFlipd(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            spatial_axis=[1],\n",
    "            prob=0.10,\n",
    "        ),\n",
    "        RandFlipd(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            spatial_axis=[2],\n",
    "            prob=0.10,\n",
    "        ),\n",
    "        RandRotate90d(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            prob=0.10,\n",
    "            max_k=3,\n",
    "        ),\n",
    "        RandShiftIntensityd(\n",
    "            keys=[\"image\"],\n",
    "            offsets=0.10,\n",
    "            prob=0.50,\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "val_transforms = Compose(\n",
    "    [\n",
    "        LoadImaged(keys=[\"image\", \"label\"], ensure_channel_first=True),\n",
    "        ScaleIntensityRanged(keys=[\"image\"], a_min=-175, a_max=250, b_min=0.0, b_max=1.0, clip=True),\n",
    "        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
    "        Spacingd(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            pixdim=(1.5, 1.5, 2.0),\n",
    "            mode=(\"bilinear\", \"nearest\"),\n",
    "        ),\n",
    "        EnsureTyped(keys=[\"image\", \"label\"], device=device, track_meta=True),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download dataset and format in the folder\n",
    "1. Download dataset from here: https://www.synapse.org/#!Synapse:syn3193805/wiki/89480. After you open the link, navigate to the \"Files\" tab, then download Abdomen/RawData.zip.\n",
    "\n",
    "    Note that you may need to register for an account on Synapse and consent to use agreements before being able to view/download this file. There are options to download directly from the browser or from the command line; please refer to Synapse API documentation for more info.\n",
    "\n",
    "\n",
    "2. After downloading the zip file, unzip. Then put images from `RawData/Training/img` in `./data/imagesTr`, and put labels from `RawData/Training/label` in `./data/labelsTr`.\n",
    "\n",
    "\n",
    "3. Make a JSON file to define train/val split and other relevant parameters. Place the JSON file at `./data/dataset_0.json`.\n",
    "\n",
    "    You can download an example of the JSON file [here](https://drive.google.com/file/d/1qcGh41p-rI3H_sQ0JwOAhNiQSXriQqGi/view?usp=sharing), or, equivalently, use the following `wget` command. If you would like to use this directly, please move it into the `./data` folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading dataset:   0%|          | 0/8 [00:06<?, ?it/s]\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "applying transform <monai.transforms.io.dictionary.LoadImaged object at 0x7f922e838f40>",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/transform.py:141\u001b[0m, in \u001b[0;36mapply_transform\u001b[0;34m(transform, data, map_items, unpack_items, log_stats, lazy, overrides)\u001b[0m\n\u001b[1;32m    140\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m [_apply_transform(transform, item, unpack_items, lazy, overrides, log_stats) \u001b[38;5;28;01mfor\u001b[39;00m item \u001b[38;5;129;01min\u001b[39;00m data]\n\u001b[0;32m--> 141\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_apply_transform\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtransform\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43munpack_items\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlazy\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moverrides\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlog_stats\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    142\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    143\u001b[0m     \u001b[38;5;66;03m# if in debug mode, don't swallow exception so that the breakpoint\u001b[39;00m\n\u001b[1;32m    144\u001b[0m     \u001b[38;5;66;03m# appears where the exception was raised.\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/transform.py:98\u001b[0m, in \u001b[0;36m_apply_transform\u001b[0;34m(transform, data, unpack_parameters, lazy, overrides, logger_name)\u001b[0m\n\u001b[1;32m     96\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m transform(\u001b[38;5;241m*\u001b[39mdata, lazy\u001b[38;5;241m=\u001b[39mlazy) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(transform, LazyTrait) \u001b[38;5;28;01melse\u001b[39;00m transform(\u001b[38;5;241m*\u001b[39mdata)\n\u001b[0;32m---> 98\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m transform(data, lazy\u001b[38;5;241m=\u001b[39mlazy) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(transform, LazyTrait) \u001b[38;5;28;01melse\u001b[39;00m \u001b[43mtransform\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/io/dictionary.py:162\u001b[0m, in \u001b[0;36mLoadImaged.__call__\u001b[0;34m(self, data, reader)\u001b[0m\n\u001b[1;32m    161\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, meta_key, meta_key_postfix \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mkey_iterator(d, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmeta_keys, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmeta_key_postfix):\n\u001b[0;32m--> 162\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_loader\u001b[49m\u001b[43m(\u001b[49m\u001b[43md\u001b[49m\u001b[43m[\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreader\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    163\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_loader\u001b[38;5;241m.\u001b[39mimage_only:\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/io/array.py:294\u001b[0m, in \u001b[0;36mLoadImage.__call__\u001b[0;34m(self, filename, reader)\u001b[0m\n\u001b[1;32m    293\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mensure_channel_first:\n\u001b[0;32m--> 294\u001b[0m     img \u001b[38;5;241m=\u001b[39m \u001b[43mEnsureChannelFirst\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimg\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    295\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mimage_only:\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/utility/array.py:200\u001b[0m, in \u001b[0;36mEnsureChannelFirst.__call__\u001b[0;34m(self, img, meta_dict)\u001b[0m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstrict_check:\n\u001b[0;32m--> 200\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(msg)\n\u001b[1;32m    201\u001b[0m warnings\u001b[38;5;241m.\u001b[39mwarn(msg)\n",
      "\u001b[0;31mValueError\u001b[0m: Metadata not available and channel_dim=None, EnsureChannelFirst is not in use.",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[25], line 7\u001b[0m\n\u001b[1;32m      5\u001b[0m datalist \u001b[38;5;241m=\u001b[39m load_decathlon_datalist(datasets, \u001b[38;5;28;01mTrue\u001b[39;00m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtraining\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      6\u001b[0m val_files \u001b[38;5;241m=\u001b[39m load_decathlon_datalist(datasets, \u001b[38;5;28;01mTrue\u001b[39;00m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvalidation\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 7\u001b[0m train_ds \u001b[38;5;241m=\u001b[39m \u001b[43mCacheDataset\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdatalist\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtransform\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrain_transforms\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcache_num\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m24\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcache_rate\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1.0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     12\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnum_workers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m8\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     13\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m     14\u001b[0m train_loader \u001b[38;5;241m=\u001b[39m ThreadDataLoader(train_ds, num_workers\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m, batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m, shuffle\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m     15\u001b[0m val_ds \u001b[38;5;241m=\u001b[39m CacheDataset(data\u001b[38;5;241m=\u001b[39mval_files, transform\u001b[38;5;241m=\u001b[39mval_transforms, cache_num\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m6\u001b[39m, cache_rate\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1.0\u001b[39m, num_workers\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m4\u001b[39m)\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/data/dataset.py:809\u001b[0m, in \u001b[0;36mCacheDataset.__init__\u001b[0;34m(self, data, transform, cache_num, cache_rate, num_workers, progress, copy_cache, as_contiguous, hash_as_key, hash_func, runtime_cache)\u001b[0m\n\u001b[1;32m    807\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_cache: \u001b[38;5;28mlist\u001b[39m \u001b[38;5;241m|\u001b[39m ListProxy \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m    808\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_hash_keys: \u001b[38;5;28mlist\u001b[39m \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m--> 809\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mset_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/data/dataset.py:836\u001b[0m, in \u001b[0;36mCacheDataset.set_data\u001b[0;34m(self, data)\u001b[0m\n\u001b[1;32m    833\u001b[0m     indices \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(\u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcache_num))\n\u001b[1;32m    835\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mruntime_cache \u001b[38;5;129;01min\u001b[39;00m (\u001b[38;5;28;01mFalse\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m):  \u001b[38;5;66;03m# prepare cache content immediately\u001b[39;00m\n\u001b[0;32m--> 836\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_cache \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_fill_cache\u001b[49m\u001b[43m(\u001b[49m\u001b[43mindices\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    837\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[1;32m    838\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mruntime_cache, \u001b[38;5;28mstr\u001b[39m) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mprocess\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mruntime_cache:\n\u001b[1;32m    839\u001b[0m     \u001b[38;5;66;03m# this must be in the main process, not in dataloader's workers\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/data/dataset.py:865\u001b[0m, in \u001b[0;36mCacheDataset._fill_cache\u001b[0;34m(self, indices)\u001b[0m\n\u001b[1;32m    863\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m ThreadPool(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_workers) \u001b[38;5;28;01mas\u001b[39;00m p:\n\u001b[1;32m    864\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprogress \u001b[38;5;129;01mand\u001b[39;00m has_tqdm:\n\u001b[0;32m--> 865\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mlist\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mtqdm\u001b[49m\u001b[43m(\u001b[49m\u001b[43mp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimap\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_load_cache_item\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mindices\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtotal\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mlen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mindices\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdesc\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mLoading dataset\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    866\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mlist\u001b[39m(p\u001b[38;5;241m.\u001b[39mimap(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_load_cache_item, indices))\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/tqdm/std.py:1181\u001b[0m, in \u001b[0;36mtqdm.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1178\u001b[0m time \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_time\n\u001b[1;32m   1180\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1181\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m obj \u001b[38;5;129;01min\u001b[39;00m iterable:\n\u001b[1;32m   1182\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m obj\n\u001b[1;32m   1183\u001b[0m         \u001b[38;5;66;03m# Update and possibly print the progressbar.\u001b[39;00m\n\u001b[1;32m   1184\u001b[0m         \u001b[38;5;66;03m# Note: does not call self.update(1) for speed optimisation.\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/multiprocessing/pool.py:870\u001b[0m, in \u001b[0;36mIMapIterator.next\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    868\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m success:\n\u001b[1;32m    869\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m value\n\u001b[0;32m--> 870\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m value\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/multiprocessing/pool.py:125\u001b[0m, in \u001b[0;36mworker\u001b[0;34m(inqueue, outqueue, initializer, initargs, maxtasks, wrap_exception)\u001b[0m\n\u001b[1;32m    123\u001b[0m job, i, func, args, kwds \u001b[38;5;241m=\u001b[39m task\n\u001b[1;32m    124\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 125\u001b[0m     result \u001b[38;5;241m=\u001b[39m (\u001b[38;5;28;01mTrue\u001b[39;00m, \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwds\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    126\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    127\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m wrap_exception \u001b[38;5;129;01mand\u001b[39;00m func \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m _helper_reraises_exception:\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/data/dataset.py:878\u001b[0m, in \u001b[0;36mCacheDataset._load_cache_item\u001b[0;34m(self, idx)\u001b[0m\n\u001b[1;32m    873\u001b[0m item \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdata[idx]\n\u001b[1;32m    875\u001b[0m first_random \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransform\u001b[38;5;241m.\u001b[39mget_index_of_first(\n\u001b[1;32m    876\u001b[0m     \u001b[38;5;28;01mlambda\u001b[39;00m t: \u001b[38;5;28misinstance\u001b[39m(t, RandomizableTrait) \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(t, Transform)\n\u001b[1;32m    877\u001b[0m )\n\u001b[0;32m--> 878\u001b[0m item \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtransform\u001b[49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfirst_random\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mthreading\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    880\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mas_contiguous:\n\u001b[1;32m    881\u001b[0m     item \u001b[38;5;241m=\u001b[39m convert_to_contiguous(item, memory_format\u001b[38;5;241m=\u001b[39mtorch\u001b[38;5;241m.\u001b[39mcontiguous_format)\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/compose.py:335\u001b[0m, in \u001b[0;36mCompose.__call__\u001b[0;34m(self, input_, start, end, threading, lazy)\u001b[0m\n\u001b[1;32m    333\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, input_, start\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m, end\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, threading\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, lazy: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m    334\u001b[0m     _lazy \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_lazy \u001b[38;5;28;01mif\u001b[39;00m lazy \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m lazy\n\u001b[0;32m--> 335\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[43mexecute_compose\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    336\u001b[0m \u001b[43m        \u001b[49m\u001b[43minput_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    337\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtransforms\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtransforms\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    338\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstart\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    339\u001b[0m \u001b[43m        \u001b[49m\u001b[43mend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    340\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmap_items\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap_items\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    341\u001b[0m \u001b[43m        \u001b[49m\u001b[43munpack_items\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43munpack_items\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    342\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlazy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_lazy\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    343\u001b[0m \u001b[43m        \u001b[49m\u001b[43moverrides\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moverrides\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    344\u001b[0m \u001b[43m        \u001b[49m\u001b[43mthreading\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mthreading\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    345\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlog_stats\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlog_stats\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    346\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    348\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/compose.py:111\u001b[0m, in \u001b[0;36mexecute_compose\u001b[0;34m(data, transforms, map_items, unpack_items, start, end, lazy, overrides, threading, log_stats)\u001b[0m\n\u001b[1;32m    109\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m threading:\n\u001b[1;32m    110\u001b[0m         _transform \u001b[38;5;241m=\u001b[39m deepcopy(_transform) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(_transform, ThreadUnsafe) \u001b[38;5;28;01melse\u001b[39;00m _transform\n\u001b[0;32m--> 111\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[43mapply_transform\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    112\u001b[0m \u001b[43m        \u001b[49m\u001b[43m_transform\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmap_items\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43munpack_items\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlazy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlazy\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moverrides\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moverrides\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlog_stats\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlog_stats\u001b[49m\n\u001b[1;32m    113\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    114\u001b[0m data \u001b[38;5;241m=\u001b[39m apply_pending_transforms(data, \u001b[38;5;28;01mNone\u001b[39;00m, overrides, logger_name\u001b[38;5;241m=\u001b[39mlog_stats)\n\u001b[1;32m    115\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m data\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/monai/transforms/transform.py:171\u001b[0m, in \u001b[0;36mapply_transform\u001b[0;34m(transform, data, map_items, unpack_items, log_stats, lazy, overrides)\u001b[0m\n\u001b[1;32m    169\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    170\u001b[0m         _log_stats(data\u001b[38;5;241m=\u001b[39mdata)\n\u001b[0;32m--> 171\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mapplying transform \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtransform\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: applying transform <monai.transforms.io.dictionary.LoadImaged object at 0x7f922e838f40>"
     ]
    }
   ],
   "source": [
    "data_dir = \"/media/chyang/data/dataset/project/registration/hipbone_raw_data/train/\"\n",
    "split_json = \"dataset.json\"\n",
    "\n",
    "datasets = data_dir + split_json\n",
    "datalist = load_decathlon_datalist(datasets, True, \"training\")\n",
    "val_files = load_decathlon_datalist(datasets, True, \"validation\")\n",
    "train_ds = CacheDataset(\n",
    "    data=datalist,\n",
    "    transform=train_transforms,\n",
    "    cache_num=24,\n",
    "    cache_rate=1.0,\n",
    "    num_workers=8,\n",
    ")\n",
    "train_loader = ThreadDataLoader(train_ds, num_workers=0, batch_size=1, shuffle=True)\n",
    "val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_num=6, cache_rate=1.0, num_workers=4)\n",
    "val_loader = ThreadDataLoader(val_ds, num_workers=0, batch_size=1)\n",
    "\n",
    "# as explained in the \"Setup transforms\" section above, we want cached training images to not have metadata, and validations to have metadata\n",
    "# the EnsureTyped transforms allow us to make this distinction\n",
    "# on the other hand, set_track_meta is a global API; doing so here makes sure subsequent transforms (i.e., random transforms for training)\n",
    "# will be carried out as Tensors, not MetaTensors\n",
    "set_track_meta(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check data shape and visualize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice_map = {\n",
    "    \"1.3.6.1.4.1.9328.50.4.0017_0000.nii.gz\": 100,\n",
    "}\n",
    "case_num = 1\n",
    "img_name = os.path.split(val_ds[case_num][\"image\"].meta[\"filename_or_obj\"])[1]\n",
    "img = val_ds[case_num][\"image\"]\n",
    "label = val_ds[case_num][\"label\"]\n",
    "img_shape = img.shape\n",
    "label_shape = label.shape\n",
    "print(f\"image shape: {img_shape}, label shape: {label_shape}\")\n",
    "plt.figure(\"image\", (18, 6))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(\"image\")\n",
    "plt.imshow(img[0, :, :, slice_map[img_name]].detach().cpu(), cmap=\"gray\")\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"label\")\n",
    "plt.imshow(label[0, :, :, slice_map[img_name]].detach().cpu())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Swin UNETR model\n",
    "\n",
    "In this section, we create a Swin UNETR model for the 14-class multi-organ segmentation. We use a feature size of 48, which is compatible with the self-supervised pre-trained weights. We also use gradient checkpointing (use_checkpoint) for more memory-efficient training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "# device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "model = SwinUNETR(\n",
    "    img_size=(96, 96, 96),\n",
    "    in_channels=1,\n",
    "    out_channels=7,\n",
    "    feature_size=48,\n",
    "    use_checkpoint=True,\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Swin UNETR encoder from self-supervised pre-trained weights\n",
    "\n",
    "In this section, we intialize the Swin UNETR encoder from pre-trained weights. The weights can be downloaded using the wget command below, or by following [this link](https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/model_swinvit.pt) to GitHub. If training from scratch is desired, please skip this section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment to download the pre-trained weights\n",
    "# !wget https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/model_swinvit.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using pretrained self-supervied Swin UNETR backbone weights !\n"
     ]
    }
   ],
   "source": [
    "weight = torch.load(\"./model_swinvit.pt\")\n",
    "model.load_from(weights=weight)\n",
    "model.float()\n",
    "# model.half()  # 将模型的参数转换为半精度浮点\n",
    "print(\"Using pretrained self-supervied Swin UNETR backbone weights !\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizer and loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.backends.cudnn.benchmark = True\n",
    "loss_function = DiceCELoss(to_onehot_y=True, softmax=True)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)\n",
    "scaler = torch.cuda.amp.GradScaler()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Execute a typical PyTorch training process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validation(epoch_iterator_val):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in epoch_iterator_val:\n",
    "            val_inputs, val_labels = (batch[\"image\"].cuda(), batch[\"label\"].cuda())\n",
    "            with torch.cuda.amp.autocast():\n",
    "                val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 4, model)\n",
    "            val_labels_list = decollate_batch(val_labels)\n",
    "            val_labels_convert = [post_label(val_label_tensor) for val_label_tensor in val_labels_list]\n",
    "            val_outputs_list = decollate_batch(val_outputs)\n",
    "            val_output_convert = [post_pred(val_pred_tensor) for val_pred_tensor in val_outputs_list]\n",
    "            dice_metric(y_pred=val_output_convert, y=val_labels_convert)\n",
    "            epoch_iterator_val.set_description(\"Validate (%d / %d Steps)\" % (global_step, 10.0))  # noqa: B038\n",
    "        mean_dice_val = dice_metric.aggregate().item()\n",
    "        dice_metric.reset()\n",
    "    return mean_dice_val\n",
    "\n",
    "\n",
    "def train(global_step, train_loader, dice_val_best, global_step_best):\n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    step = 0\n",
    "    epoch_iterator = tqdm(train_loader, desc=\"Training (X / X Steps) (loss=X.X)\", dynamic_ncols=True)\n",
    "    for step, batch in enumerate(epoch_iterator):\n",
    "        step += 1\n",
    "        x, y = (batch[\"image\"].to(device).float(), batch[\"label\"].to(device).float())\n",
    "        with torch.cuda.amp.autocast():\n",
    "            logit_map = model(x)\n",
    "            loss = loss_function(logit_map, y)\n",
    "        scaler.scale(loss).backward()\n",
    "        epoch_loss += loss.item()\n",
    "        scaler.unscale_(optimizer)\n",
    "        scaler.step(optimizer)\n",
    "        scaler.update()\n",
    "        optimizer.zero_grad()\n",
    "        epoch_iterator.set_description(  # noqa: B038\n",
    "            f\"Training ({global_step} / {max_iterations} Steps) (loss={loss:2.5f})\"\n",
    "        )\n",
    "        if (global_step % eval_num == 0 and global_step != 0) or global_step == max_iterations:\n",
    "            epoch_iterator_val = tqdm(val_loader, desc=\"Validate (X / X Steps) (dice=X.X)\", dynamic_ncols=True)\n",
    "            dice_val = validation(epoch_iterator_val)\n",
    "            epoch_loss /= step\n",
    "            epoch_loss_values.append(epoch_loss)\n",
    "            metric_values.append(dice_val)\n",
    "            if dice_val > dice_val_best:\n",
    "                dice_val_best = dice_val\n",
    "                global_step_best = global_step\n",
    "                torch.save(model.state_dict(), os.path.join(root_dir, \"best_metric_model.pth\"))\n",
    "                print(\n",
    "                    \"Model Was Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}\".format(dice_val_best, dice_val)\n",
    "                )\n",
    "            else:\n",
    "                print(\n",
    "                    \"Model Was Not Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}\".format(\n",
    "                        dice_val_best, dice_val\n",
    "                    )\n",
    "                )\n",
    "        global_step += 1\n",
    "    return global_step, dice_val_best, global_step_best"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training (7 / 30000 Steps) (loss=2.42775): 100%|██████████| 8/8 [03:25<00:00, 25.74s/it]\n",
      "Training (15 / 30000 Steps) (loss=2.17743): 100%|██████████| 8/8 [03:28<00:00, 26.04s/it]\n",
      "Training (23 / 30000 Steps) (loss=2.11350): 100%|██████████| 8/8 [03:30<00:00, 26.27s/it]\n",
      "Training (31 / 30000 Steps) (loss=1.94048): 100%|██████████| 8/8 [03:29<00:00, 26.20s/it]\n",
      "Training (39 / 30000 Steps) (loss=1.86852): 100%|██████████| 8/8 [03:28<00:00, 26.06s/it]\n",
      "Training (47 / 30000 Steps) (loss=1.88652): 100%|██████████| 8/8 [03:30<00:00, 26.25s/it]\n",
      "Training (55 / 30000 Steps) (loss=1.79510): 100%|██████████| 8/8 [03:28<00:00, 26.03s/it]\n",
      "Training (63 / 30000 Steps) (loss=1.71695): 100%|██████████| 8/8 [03:32<00:00, 26.61s/it]\n",
      "Training (71 / 30000 Steps) (loss=1.77054): 100%|██████████| 8/8 [03:32<00:00, 26.50s/it]\n",
      "Training (79 / 30000 Steps) (loss=1.65344): 100%|██████████| 8/8 [03:29<00:00, 26.21s/it]\n",
      "Training (87 / 30000 Steps) (loss=1.70147): 100%|██████████| 8/8 [03:31<00:00, 26.43s/it]\n",
      "Training (93 / 30000 Steps) (loss=1.61028):  75%|███████▌  | 6/8 [02:59<00:59, 29.99s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[31], line 12\u001b[0m\n\u001b[1;32m     10\u001b[0m metric_values \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m global_step \u001b[38;5;241m<\u001b[39m max_iterations:\n\u001b[0;32m---> 12\u001b[0m     global_step, dice_val_best, global_step_best \u001b[38;5;241m=\u001b[39m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mglobal_step\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdice_val_best\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mglobal_step_best\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     13\u001b[0m model\u001b[38;5;241m.\u001b[39mload_state_dict(torch\u001b[38;5;241m.\u001b[39mload(os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(root_dir, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbest_metric_model.pth\u001b[39m\u001b[38;5;124m\"\u001b[39m)))\n",
      "Cell \u001b[0;32mIn[30], line 30\u001b[0m, in \u001b[0;36mtrain\u001b[0;34m(global_step, train_loader, dice_val_best, global_step_best)\u001b[0m\n\u001b[1;32m     28\u001b[0m     logit_map \u001b[38;5;241m=\u001b[39m model(x)\n\u001b[1;32m     29\u001b[0m     loss \u001b[38;5;241m=\u001b[39m loss_function(logit_map, y)\n\u001b[0;32m---> 30\u001b[0m \u001b[43mscaler\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscale\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     31\u001b[0m epoch_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mitem()\n\u001b[1;32m     32\u001b[0m scaler\u001b[38;5;241m.\u001b[39munscale_(optimizer)\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/_tensor.py:525\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    515\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    517\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    518\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    523\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    524\u001b[0m     )\n\u001b[0;32m--> 525\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    526\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[1;32m    527\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/autograd/__init__.py:267\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    262\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    264\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    266\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 267\u001b[0m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    268\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    269\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    270\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    271\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    272\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    273\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    274\u001b[0m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    275\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/autograd/graph.py:744\u001b[0m, in \u001b[0;36m_engine_run_backward\u001b[0;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[1;32m    742\u001b[0m     unregister_hooks \u001b[38;5;241m=\u001b[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[1;32m    743\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 744\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    745\u001b[0m \u001b[43m        \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m    746\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[1;32m    747\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    748\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/torch/utils/checkpoint.py:1108\u001b[0m, in \u001b[0;36m_checkpoint_hook.__init__.<locals>.unpack_hook\u001b[0;34m(holder)\u001b[0m\n\u001b[1;32m   1105\u001b[0m             frame\u001b[38;5;241m.\u001b[39mx_metadatas\u001b[38;5;241m.\u001b[39mappend(frame\u001b[38;5;241m.\u001b[39mmetadata_fn(x))\n\u001b[1;32m   1106\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m holder\n\u001b[0;32m-> 1108\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21munpack_hook\u001b[39m(holder):\n\u001b[1;32m   1109\u001b[0m     gid \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39m_C\u001b[38;5;241m.\u001b[39m_current_graph_task_id()\n\u001b[1;32m   1110\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m gid \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m   1111\u001b[0m         \u001b[38;5;66;03m# generate a temporary id if we trigger unpack outside of a backward call\u001b[39;00m\n",
      "File \u001b[0;32m_pydevd_bundle/pydevd_cython.pyx:1457\u001b[0m, in \u001b[0;36m_pydevd_bundle.pydevd_cython.SafeCallWrapper.__call__\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m_pydevd_bundle/pydevd_cython.pyx:1758\u001b[0m, in \u001b[0;36m_pydevd_bundle.pydevd_cython.ThreadTracer.__call__\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/anaconda3/envs/monai/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_is_thread_alive.py:9\u001b[0m, in \u001b[0;36mis_thread_alive\u001b[0;34m(t)\u001b[0m\n\u001b[1;32m      6\u001b[0m _temp \u001b[38;5;241m=\u001b[39m threading\u001b[38;5;241m.\u001b[39mThread()\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(_temp, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m_is_stopped\u001b[39m\u001b[38;5;124m'\u001b[39m):  \u001b[38;5;66;03m# Python 3.x has this\u001b[39;00m\n\u001b[0;32m----> 9\u001b[0m     \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mis_thread_alive\u001b[39m(t):\n\u001b[1;32m     10\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m t\u001b[38;5;241m.\u001b[39m_is_stopped\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(_temp, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m_Thread__stopped\u001b[39m\u001b[38;5;124m'\u001b[39m):  \u001b[38;5;66;03m# Python 2.x has this\u001b[39;00m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "max_iterations = 30000\n",
    "eval_num = 500\n",
    "post_label = AsDiscrete(to_onehot=14)\n",
    "post_pred = AsDiscrete(argmax=True, to_onehot=14)\n",
    "dice_metric = DiceMetric(include_background=True, reduction=\"mean\", get_not_nans=False)\n",
    "global_step = 0\n",
    "dice_val_best = 0.0\n",
    "global_step_best = 0\n",
    "epoch_loss_values = []\n",
    "metric_values = []\n",
    "while global_step < max_iterations:\n",
    "    global_step, dice_val_best, global_step_best = train(global_step, train_loader, dice_val_best, global_step_best)\n",
    "model.load_state_dict(torch.load(os.path.join(root_dir, \"best_metric_model.pth\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"train completed, best_metric: {dice_val_best:.4f} \" f\"at iteration: {global_step_best}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Plot the loss and metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(\"train\", (12, 6))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(\"Iteration Average Loss\")\n",
    "x = [eval_num * (i + 1) for i in range(len(epoch_loss_values))]\n",
    "y = epoch_loss_values\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.plot(x, y)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Val Mean Dice\")\n",
    "x = [eval_num * (i + 1) for i in range(len(metric_values))]\n",
    "y = metric_values\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.plot(x, y)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "monai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
